from nt import write
import torch
from torch.utils.data import dataset, DataLoader
from torch.utils.tensorboard.writer import SummaryWriter
import torchvision
from torch import nn

from nn_module import sheamus

dataset = torchvision.datasets.CIFAR10(root='./data', train=False, transform=torchvision.transforms.ToTensor(), download=True)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True)

class Sheamus(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=3, stride=1, padding=0)
        
    
    def forward(self, x):
        x = self.conv1(x)
        return x

sheamus = Sheamus()
print(sheamus)

write = SummaryWriter(log_dir='logs')
step = 0

for data in dataloader:
    imgs, targets = data
    print(imgs.shape)
    output = sheamus(imgs)
    print(output.shape)
    write.add_images('input', imgs, global_step=step)
    output = torch.reshape(output, (-1, 3, output.shape[2], output.shape[3]))
    write.add_images('output', output, global_step=step)
    step += 1

write.close()